
#include "mat_mul_sub_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include "tiling/tiling_api.h"
using namespace matmul_tiling;

namespace optiling {
    constexpr uint32_t REPEAT_NUM = 256;
    constexpr uint32_t BUFFER_NUM = 2;  

    uint32_t RoundUp(uint32_t a, uint32_t b)
    { 
        if(b == 0) {
            return 0;
        }
        return ((a + b - 1) / b) * b;
    }      
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto socVersion = ascendcPlatform.GetSocVersion();
    auto aivcoreNum = ascendcPlatform.GetCoreNumAiv();
   
    std::cout << "aivcoreNum = " << aivcoreNum <<std::endl;

    auto shape_a = context->GetInputTensor(0)->GetOriginShape();
    auto shape_b = context->GetInputTensor(1)->GetOriginShape();
    int32_t numshapes = context->GetInputTensor(2)->GetOriginShape().GetDimNum();

    int32_t M = shape_a.GetDim(0);
    int32_t N = shape_b.GetDim(1);
    int32_t K = shape_a.GetDim(1);
    
    int32_t baseM; 
    int32_t baseN;

    std::cout << "M = " << M <<std::endl;
    std::cout << "N = " << N <<std::endl;
    std::cout << "H = " << K <<std::endl;

    uint32_t dTypeSize;
    auto dataType = context->GetInputTensor(0)->GetDataType();  
    switch (dataType) {
        case ge::DT_FLOAT:
            dTypeSize = 4;
            std::cout << "This is float!!!" << std::endl;
        break;
        case ge::DT_FLOAT16:
            dTypeSize = 2;
            std::cout << "This is float16!!!" << std::endl;
        break;           
        default:
            std::cout << "The input shape is error!!!" << std::endl;
            return ge::GRAPH_FAILED;
        break;
    }    

    MultiCoreMatmulTiling cubeTiling(ascendcPlatform);
 
    if (dataType == ge::DT_FLOAT) {
        cubeTiling.SetAType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);
        cubeTiling.SetBType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);
        cubeTiling.SetCType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);
        cubeTiling.SetBiasType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);

    }
    else {
        cubeTiling.SetAType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16);
        cubeTiling.SetBType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16);   
        cubeTiling.SetCType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16); 
        cubeTiling.SetBiasType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16);           
    }

    cubeTiling.SetShape(M, N, K);
    cubeTiling.SetOrgShape(M, N, K);
   
    if (M <=32) {
        cubeTiling.SetDim(2);
        baseM = 32; 
        baseN = 16; 
        cubeTiling.SetFixSplit(baseM, baseN, -1);        
    }
    else
    {     
        cubeTiling.SetDim(aivcoreNum);  
        cubeTiling.SetFixSplit(-1, -1, -1);
    }
    // 
    cubeTiling.SetBias(true);
    cubeTiling.SetBufferSpace(-1, -1, -1);


    MatMulSubTilingData tiling;
    if (cubeTiling.GetTiling(tiling.cubeTilingData) == -1) {
        std::cout << "cubeTiling.GetTiling error!!!!" <<std::endl;
        return ge::GRAPH_FAILED;
    }

    int32_t dim, mDim, nDim;
    int ret1 = cubeTiling.GetCoreNum(dim, mDim, nDim);

    int calc_baseM = cubeTiling.GetBaseM(); 
    int calc_baseN = cubeTiling.GetBaseN(); 
    int calc_baseK = cubeTiling.GetBaseK(); 

    int32_t calc_singleM, calc_singleN, calc_singleK;
    int ret2 = cubeTiling.GetSingleShape(calc_singleM, calc_singleN, calc_singleK);

    uint32_t coreNum = dim / 2;
    context->SetBlockDim((coreNum));

    uint64_t ub_size;
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);    
    uint32_t totalLength = M * N;
    std::cout << "totalLength = " << totalLength <<std::endl;     
    uint32_t ub_tileLength = RoundUp((ub_size / BUFFER_NUM / dTypeSize), (REPEAT_NUM / dTypeSize));
    std::cout << "ub_tileLength = " << ub_tileLength <<std::endl;     
    uint32_t sigleCoreNum = (totalLength + dim - 1) / dim;
    std::cout << "sigleCoreNum = " << sigleCoreNum <<std::endl;   
    uint32_t tileNum = (sigleCoreNum + ub_tileLength - 1) / ub_tileLength;
    std::cout << "tileNum = " << tileNum <<std::endl; 
    
    uint32_t tileLength = (sigleCoreNum + tileNum - 1)/ tileNum;
    
    tileLength = RoundUp(tileLength, ((REPEAT_NUM / dTypeSize)));
    
    if(tileLength > totalLength) {
        tileLength = totalLength;
    }
    std::cout << "tileLength = " << tileLength <<std::endl; 
    
    int32_t x3Dims = context->GetInputTensor(2)->GetOriginShape().GetDimNum();
    if(x3Dims == 2) {
        context->SetTilingKey(1);
        std::cout << "context->SetTilingKey(1) "  <<std::endl; 
    }
    else {
        context->SetTilingKey(2);
        std::cout << "context->SetTilingKey(2) "  <<std::endl; 
        totalLength = M;
        tileLength = 1;
    }

    tiling.set_totalLength(totalLength);
    tiling.set_tileLength(tileLength);
    
    std::cout << "totalLength = " << totalLength <<std::endl; 
    std::cout << "tileLength = " << tileLength <<std::endl;

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t userWorkspaceSize;
    userWorkspaceSize = 0;
    
    std::cout << "userWorkspaceSize = " << userWorkspaceSize <<std::endl; 

    size_t systemWorkspaceSize = static_cast<size_t>(ascendcPlatform.GetLibApiWorkSpaceSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = userWorkspaceSize + systemWorkspaceSize;
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class MatMulSub : public OpDef {
public:
    explicit MatMulSub(const char* name) : OpDef(name)
    {
        this->Input("x1")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x2")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x3")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(MatMulSub);
}
